package bintree.heap.leetcode;

import bintree.heap.PriorityQueue;

import java.util.Arrays;
import java.util.Map;
import java.util.TreeMap;

/**
 * 前k个高频元素
 * https://leetcode-cn.com/problems/top-k-frequent-elements/
 * @author yuisama
 * @date 2021/12/13 16:37
 **/
public class Num347 {
    // 存储元素以及出现频次
    private class Freq implements Comparable<Freq>{
        int val;
        int times;

        public Freq(int val, int times) {
            this.val = val;
            this.times = times;
        }

        @Override
        public int compareTo(Freq o) {
            if (this.times < o.times) return -1;
            else if (this.times > o.times) return 1;
            return 0;
        }
    }
    public int[] topKFrequent(int[] nums, int k) {
        // 先用TreeMap存储每个元素以及出现频次
        Map<Integer,Integer> map = new TreeMap<>();
        for(int i : nums) {
            if (map.containsKey(i)) {
                map.put(i,map.get(i) + 1);
            }else {
                map.put(i,1);
            }
        }
        // 使用优先级队列(最小堆)存储前K个键值对
        PriorityQueue<Freq> priorityQueue = new PriorityQueue<>();
        for (int key : map.keySet()) {
            if (priorityQueue.getSize() < k) {
                priorityQueue.enqueue(new Freq(key,map.get(key)));
            }else {
                if (map.get(key) > priorityQueue.getFont().times) {
                    priorityQueue.dequeue();
                    priorityQueue.enqueue(new Freq(key,map.get(key)));
                }
            }
        }
        // 最终存储前k个高频元素的数组
        int[] ret = new int[k];
        int i = 0;
        while (!priorityQueue.isEmpty()) {
            ret[i++] = priorityQueue.dequeue().val;
        }
        return ret;
    }

    public static void main(String[] args) {
        int[] test = new int[] {4,1,-1,2,-1,2,3};
        int[] ret = new Num347().topKFrequent(test,2);
        System.out.println(Arrays.toString(ret));
    }
}